package org.ala.tiktools.sharddb.tx;


import java.sql.Connection;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import org.ala.tiktools.logs.LoggerFactory;
import org.ala.tiktools.tools.list.Lists;
import org.slf4j.Logger;
import org.springframework.jdbc.datasource.ConnectionHolder;
import org.springframework.jdbc.datasource.DataSourceTransactionManagerExpander;
import org.springframework.jdbc.datasource.DataSourceUtils;
import org.springframework.lang.Nullable;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionException;
import org.springframework.transaction.support.DefaultTransactionStatus;
import org.springframework.transaction.support.TransactionSynchronization;
import org.springframework.transaction.support.TransactionSynchronizationManager;


/**
 * 分片数据源事务管理器
 * <p>	就是因为 DataSourceTransactionObject 和 rollback 是private的，费了多少事。。。
 *
 * @author ala
 * @date 2024-12-30 09:30
 */
public class ShardDataSourceTransactionManager extends DataSourceTransactionManagerExpander {
    private static final long serialVersionUID = 1L;
    static Logger log = LoggerFactory.getLogger("shard", "transaction");
    
	/**
     *  执行预提交的线程池
     *  <p> 执行事务的conn是被setAutoCommit的，而这个conn是强绑定线程环境的
     *  <p> 所以这里只能用线程池去执行预提交
     */
    protected ThreadPoolExecutor threadPool;
    protected int coreSize = 1, maxSize = 5, queueSize = 1024;

    /**
     * 默认的事务传播方式
     * <p>  融入
     */
    static ShardTransactionDefinition DEFAULT_TRANSACTION_DEFINITION;
    static {
        //	设置本地事务的传播性：挂起原有事务，新开事务
    	ShardTransactionDefinition ed = new ShardTransactionDefinition();
        ed.setPropagationBehavior(TransactionDefinition.PROPAGATION_REQUIRED);
        DEFAULT_TRANSACTION_DEFINITION = ed;
    }

    public void afterPropertiesSet() {
        super.afterPropertiesSet();

        //  初始化线程池
        initThreadPool();
    }
    protected void initThreadPool() {
        AtomicInteger cnt = new AtomicInteger();
        threadPool = new ThreadPoolExecutor(
                coreSize, maxSize, 180, TimeUnit.SECONDS,
                new ArrayBlockingQueue<Runnable>(queueSize),
                new ThreadFactory() {
                    @Override
                    public Thread newThread(Runnable r) {
                        Thread t = new Thread(r);
                        t.setDaemon(false);
                        t.setName(String.format("ShardTransactionThread_%d", cnt.getAndIncrement()));
                        return t;
                    }
                });
    }

    /**
     *  开事务
     *  <p>	这里开的是事务协调者
     *  <p>	事务协调者在执行预提交代码之前，才会拿到真实的TransactionStatus
     */
    public IShardTransaction open(long timeout) {
        return open(timeout, DEFAULT_TRANSACTION_DEFINITION);
    }
    /**
     *  开事务
     *  <p>	这里开的是事务协调者
     *  <p>	事务协调者在执行预提交代码之前，才会拿到真实的TransactionStatus
     */
    public IShardTransaction open(long timeout, ThreadPoolExecutor threadPool) {
        return open(timeout, DEFAULT_TRANSACTION_DEFINITION, threadPool);
    }
    /**
     *  开事务
     *  <p>	这里开的是事务协调者
     *  <p>	事务协调者在执行预提交代码之前，才会拿到真实的TransactionStatus
     */
    public IShardTransaction open(long timeout, TransactionDefinition td) {
        return open(timeout, DEFAULT_TRANSACTION_DEFINITION, threadPool);
    }
    /**
     *  开事务
     *  <p>	这里开的是事务协调者
     *  <p>	事务协调者在执行预提交代码之前，才会拿到真实的TransactionStatus
     */
    public IShardTransaction open(long timeout, TransactionDefinition td, ThreadPoolExecutor threadPool) {
        ShardTransaction stx = new ShardTransaction();
        stx.threadPool = threadPool;
        stx.transactionDefinition = td;
        stx.transactionManager = this;
        stx.timeout = timeout;
        return stx;
    }
    
    
	
	@Override
	protected Object doGetTransaction(TransactionDefinition definition) {
		DataSourceTransactionObject txObject = null;
		//	判断是独立事务，开始分布式事务
		if (definition instanceof ShardTransactionDefinition) {
			txObject = new ShardDataSourceTransactionObject();
		} else {
			txObject = new DataSourceTransactionObject();
		}
		
		txObject.setSavepointAllowed(isNestedTransactionAllowed());
		ConnectionHolder conHolder =
				(ConnectionHolder) TransactionSynchronizationManager.getResource(obtainDataSource());
		txObject.setConnectionHolder(conHolder, false);
		return txObject;
	}
	
	
	@Override
	protected DefaultTransactionStatus newTransactionStatus(TransactionDefinition definition, Object transaction, boolean newTransaction, boolean newSynchronization, boolean debug, Object suspendedResources) {
		boolean actualNewSynchronization = newSynchronization && !TransactionSynchronizationManager.isSynchronizationActive();
		if (definition instanceof ShardTransactionDefinition) {
			return new ShardTransactionStatus(transaction, newTransaction, newSynchronization, actualNewSynchronization, debug, suspendedResources);
		} else {
			return new DefaultTransactionStatus(transaction, newTransaction, actualNewSynchronization, definition.isReadOnly(), debug, suspendedResources);
		}
	}
	
	
	@Override
	protected void prepareSynchronization(DefaultTransactionStatus status, TransactionDefinition definition) {
		super.prepareSynchronization(status, definition);
		
		if (status instanceof ShardTransactionStatus) {
			((ShardTransactionStatus) status).setDataSource(obtainDataSource());
		}
	}
	
	
	@Override
	protected void doCleanupAfterCompletion(Object transaction) {
		DataSourceTransactionObject txObject = (DataSourceTransactionObject) transaction;

		//	只有非分布式事务，才需要清除线程环境
		//	分布式事务的线程环境在预提交完成后就清除了
		if (txObject.isNewConnectionHolder() && !(txObject instanceof ShardDataSourceTransactionObject)) {
			TransactionSynchronizationManager.unbindResource(obtainDataSource());
		}

		// Reset connection.
		Connection con = txObject.getConnectionHolder().getConnection();
		try {
			if (txObject.isMustRestoreAutoCommit()) {
				con.setAutoCommit(true);
			}
			DataSourceUtils.resetConnectionAfterTransaction(
					con, txObject.getPreviousIsolationLevel(), txObject.isReadOnly());
		}
		catch (Throwable ex) {
			log.debug("Could not reset JDBC Connection after transaction", ex);
		}

		if (txObject.isNewConnectionHolder()) {
			if (log.isDebugEnabled()) {
				log.debug("Releasing JDBC Connection [" + con + "] after transaction");
			}
			DataSourceUtils.releaseConnection(con, this.dataSource);
		}

		txObject.getConnectionHolder().clear();
	}


	@Override
	protected void triggerBeforeCommit(DefaultTransactionStatus status) {
		if (status instanceof ShardTransactionStatus) {
			((ShardDataSourceTransactionObject) status.getTransaction()).triggerBeforeCommit(status.isReadOnly());
		} else {
			super.triggerBeforeCommit(status);
		}
	}
	@Override
	protected void triggerBeforeCompletion(DefaultTransactionStatus status) {
		if (status instanceof ShardTransactionStatus) {
			((ShardDataSourceTransactionObject) status.getTransaction()).triggerBeforeCompletion();
		} else {
			super.triggerBeforeCompletion(status);
		}
	}
	@Override
	protected void triggerAfterCommit(DefaultTransactionStatus status) {
		if (status instanceof ShardTransactionStatus) {
			((ShardDataSourceTransactionObject) status.getTransaction()).triggerAfterCommit();
		} else {
			super.triggerAfterCommit(status);
		}
	}
	@Override
	protected void triggerAfterCompletion(DefaultTransactionStatus status, int completionStatus) {
		if (status instanceof ShardTransactionStatus) {
			((ShardDataSourceTransactionObject) status.getTransaction()).triggerAfterCompletion(status, completionStatus);
		} else {
			super.triggerAfterCompletion(status, completionStatus);
		}
	}


	/**
	 * 持有ConnectionHolder
	 */
	public static class ShardDataSourceTransactionObject extends DataSourceTransactionObject {
		/**
		 * 	同步任务
		 * 	<p>	单独拿出来保存，因为线程环境预提交完成后就清掉了
		 */
		protected List<TransactionSynchronization> synchronizations;

		public List<TransactionSynchronization> getSynchronizations() {
			return synchronizations;
		}
		public void setSynchronizations(List<TransactionSynchronization> synchronizations) {
			this.synchronizations = synchronizations;
		}
		public ShardDataSourceTransactionObject appendSynchronizations(TransactionSynchronization ts) {
			if (synchronizations == null) {
				synchronizations = new ArrayList<>();
			}
			synchronizations.add(ts);
			return this;
		}
		
		
		/**
		 * 代替原线程环境中的 synchronizations
		 */
		public void triggerBeforeCommit(boolean readOnly) {
			if (!Lists.isEmpty(synchronizations)) {
				for (TransactionSynchronization synchronization : synchronizations) {
					synchronization.beforeCommit(readOnly);
				}
			}
		}
		public void triggerBeforeCompletion() {
			if (!Lists.isEmpty(synchronizations)) {
				for (TransactionSynchronization synchronization : synchronizations) {
					try {
						synchronization.beforeCompletion();
					}
					catch (Throwable tsex) {
						log.error("TransactionSynchronization.beforeCompletion threw exception", tsex);
					}
				}
			}
		}
		public void triggerAfterCommit() {
			invokeAfterCommit(synchronizations);
		}
		public void invokeAfterCommit(@Nullable List<TransactionSynchronization> synchronizations) {
			if (!Lists.isEmpty(synchronizations)) {
				for (TransactionSynchronization synchronization : synchronizations) {
					synchronization.afterCommit();
				}
			}
		}
		protected void invokeAfterCompletion(List<TransactionSynchronization> synchronizations, int completionStatus) {
			if (!Lists.isEmpty(synchronizations)) {
				for (TransactionSynchronization synchronization : synchronizations) {
					try {
						synchronization.afterCompletion(completionStatus);
					}
					catch (Throwable tsex) {
						log.error("TransactionSynchronization.afterCompletion threw exception", tsex);
					}
				}
			}
		}
		protected void triggerAfterCompletion(DefaultTransactionStatus status, int completionStatus) {
			if (status.isNewSynchronization()) {
				if (!status.hasTransaction() || status.isNewTransaction()) {
					if (status.isDebug()) {
						log.trace("Triggering afterCompletion synchronization");
					}
					invokeAfterCompletion(synchronizations, completionStatus);
//					synchronizations.clear();
				}
				else if (!synchronizations.isEmpty()) {
					registerAfterCompletionWithExistingTransaction(status.getTransaction(), synchronizations);
				}
			}
		}
		protected void registerAfterCompletionWithExistingTransaction(Object transaction, List<TransactionSynchronization> synchronizations) throws TransactionException {
			log.debug("Cannot register Spring after-completion synchronization with existing transaction - processing Spring after-completion callbacks immediately, with outcome status 'unknown'");
			invokeAfterCompletion(synchronizations, TransactionSynchronization.STATUS_UNKNOWN);
		}
	}
    
}
